package com.buaa.pm.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.buaa.pm.entity.ExtractRecord;
import com.buaa.pm.mapper.ExtractRecordMapper;
import com.buaa.pm.model.vo.ExtractRecordVO;
import com.buaa.pm.service.ExtractRecordService;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

@Service
public class ExtractRecordServiceImpl extends ServiceImpl<ExtractRecordMapper, ExtractRecord> implements ExtractRecordService {
    @Override
    public List<ExtractRecordVO> listExtractRecordByTask(String taskId) {
        List<ExtractRecordVO> extractRecordVOList = new ArrayList<>();

        QueryWrapper<ExtractRecord> queryWrapper = new QueryWrapper<>();
        queryWrapper.lambda().eq(ExtractRecord::getTaskId, taskId).orderByDesc(ExtractRecord::getContent);
        List<ExtractRecord> extractRecords = list(queryWrapper);

        Map<String, Map<String, Object>> recordsMap = new HashMap<>();
        // 初始化数据结构
        for (ExtractRecord extractRecord : extractRecords) {
            String paragraphId = extractRecord.getParagraphId();
            String tag = extractRecord.getTag();
            if (tag.equals("Metal_Source") || tag.equals("Solvent") || tag.equals("Organic_Linker") || tag.equals("Modulator")) {
                recordsMap.computeIfAbsent(paragraphId, k -> new HashMap<>())
                        .computeIfAbsent(tag, k -> new ArrayList<>());
            } else {
                recordsMap.computeIfAbsent(paragraphId, k -> new HashMap<>())
                        .computeIfAbsent(tag, k -> new HashSet<>());
            }
        }
        // 批量查询amount对应的ExtractRecord
        Set<String> amountIds = extractRecords.stream()
                .filter(extractRecord -> extractRecord.getAmount() != null)
                .map(ExtractRecord::getAmount)
                .collect(Collectors.toSet());

        Map<String, ExtractRecord> amountMap;
        if (!amountIds.isEmpty()) {
            QueryWrapper<ExtractRecord> amountQueryWrapper = new QueryWrapper<>();
            amountQueryWrapper.lambda().in(ExtractRecord::getId, amountIds);
            List<ExtractRecord> amountRecords = list(amountQueryWrapper);
            amountMap = amountRecords.stream()
                    .collect(Collectors.toMap(ExtractRecord::getId, Function.identity()));
        } else {
            amountMap = new HashMap<>();
        }

        // 并行处理数据
        extractRecords.parallelStream().forEach(extractRecord -> {
            String paragraphId = extractRecord.getParagraphId();
            String tag = extractRecord.getTag();
            String content = extractRecord.getContent();
            String amount = extractRecord.getAmount();

            if (tag.equals("Metal_Source") || tag.equals("Solvent") || tag.equals("Organic_Linker") || tag.equals("Modulator")) {
                if (amount != null) {
                    ExtractRecord amountRecord = amountMap.get(amount);
                    if (amountRecord != null) {
                        Map<String, String> data = new HashMap<>();
                        data.put("precursor_name", content);
                        data.put("amount", amountRecord.getContent());
                        List<Map<String, String>> dataList = (List<Map<String, String>>) recordsMap.get(paragraphId).get(tag);
                        synchronized (dataList) {
                            dataList.add(data);
                        }
                    }
                }
            } else {
                HashSet<String> dataSet = (HashSet<String>) recordsMap.get(paragraphId).get(tag);
                synchronized (dataSet) {
                    dataSet.add(content);
                }
            }
        });

        // 构建结果对象
        for (Map.Entry<String, Map<String, Object>> entry : recordsMap.entrySet()) {
            String paragraphId = entry.getKey();
            Map<String, Object> mofRecords = entry.getValue();

            ExtractRecordVO extractRecordVO = new ExtractRecordVO();
            extractRecordVO.setParagraphId(paragraphId);
            extractRecordVO.setParam(mofRecords);

            extractRecordVOList.add(extractRecordVO);
        }
        return extractRecordVOList;
    }

    @Override
    public List<ExtractRecord> listExtractRecordByPara(String paragraphId,String taskId){
        QueryWrapper<ExtractRecord> queryWrapper = new QueryWrapper<>();
        queryWrapper.lambda().eq(ExtractRecord::getParagraphId, paragraphId).eq(ExtractRecord::getTaskId, taskId)
                .orderByDesc(ExtractRecord::getTag);
        return list(queryWrapper);
    }
}
